import json
import os
from locust import HttpUser, TaskSet, task
from jsonpath_rw import parser

path = os.environ['LOCUSTPATH']
method = os.environ['LOCUSTMETHOD']
data = json.loads(os.environ['LOCUSTREQUEST'])
header = {"content-type": f"{os.environ['LOCUSTHEADER']}"}


class LocustTest(TaskSet):
    @task(1)
    def inner(self):
        url = path
        if method.lower() == 'get':
            with self.client.get(url, params=data, catch_response=True) as res:
                if res.status_code != 200:
                    print(res.json())
                    res.failure("Failed!")
                else:
                    print(res.json())
                    status = parse_resp_json(res.json(), 'status')
                    res.success() if status == 1 else res.failure('Failed!')
        elif method.lower() == 'post':
            with self.client.post(url, json=json.dumps(data), catch_response=True) as res:
                if res.status_code != 200:
                    res.failure("Failed!")
                else:
                    status = parse_resp_json(res.json(), 'status')
                    res.success() if status == 1 else res.failure('Failed!')


class WebSitUser(HttpUser):
    tasks = [LocustTest]
    min_wait = os.environ['MINWAIT']
    max_wait = os.environ['MAXWAIT']


def parse_resp_json(json_data, json_path, index=0):
    json_path_expr = parser.parse(json_path)
    if isinstance(json_data, list):
        results = [match.value for match in json_path_expr.find(
            json_data[index])]
    else:
        results = [match.value for match in json_path_expr.find(json_data)]
    if not results:
        raise KeyError('Key not found!')
    if len(results) == 1:
        return results[0]
    else:
        return results
